#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <cuda.h>
#include "common.h"

// commented, so we can specify the number of threads
//#define NUM_THREADS 256
#define P_PER_B 32

#define IDX(i,j,NB) ((i)*(NB)+(j))
#define OUTBOUND(i,j,NB) ((i) < 0 || (i) >= (NB) || (j) < 0 || (j) >= (NB))


extern double size;
//
//  benchmarking program
//

__device__ void apply_force_gpu(particle_t &particle, particle_t &neighbor)
{
    double dx = neighbor.x - particle.x;
    double dy = neighbor.y - particle.y;
    double r2 = dx * dx + dy * dy;
    if( r2 > cutoff*cutoff )
        return;
    r2 = (r2 > min_r*min_r) ? r2 : min_r*min_r;
    double coef = ( 1 - cutoff / sqrt(r2) ) / r2 / mass;
    particle.ax += coef * dx;
    particle.ay += coef * dy;

}

__device__ int getBinID(particle_t & particle, int NB)
{
    return floor(particle.x / cutoff) * NB + floor(particle.y / cutoff); 
}

__global__ void assign_bins_gpu(particle_t * particles, int n, unsigned int * bins, unsigned int * binMax, int NB)
{
    // Get thread (particle) ID
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if(tid >= n) return;

    int binid = getBinID(particles[tid],NB);
    int position = atomicInc(binMax+binid, P_PER_B);
    bins[binid*P_PER_B + position] = tid;
}

__global__ void compute_forces_gpu(particle_t * particles, int n, unsigned int * bins, unsigned int * binMax, int NB)
{
    // Get thread (particle) ID
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if(tid >= n) return;
    // use local cache
    particle_t pt = particles[tid];

    int dstbinid = getBinID(pt, NB);
    int i = dstbinid / NB;
    int j = dstbinid % NB;
    int istart = max(i-1,0), iend = min(i+2,NB),
        jstart = max(j-1,0), jend = min(j+2,NB);
    for (int _i = istart; _i < iend; _i ++)
    { 
        for (int _j = jstart; _j < jend; _j ++)
        {
            unsigned int * neighborbin = & bins[IDX(_i,_j,NB) * P_PER_B];
            int pnum = binMax[IDX(_i,_j,NB)];
            for (int src = 0; src < pnum; src ++)
            { 
                apply_force_gpu(pt, particles[neighborbin[src]]);
            } 
        }
    }
    particles[tid] = pt;

}

__global__ void move_gpu (particle_t * particles, int n, double size)
{

    // Get thread (particle) ID
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if(tid >= n) return;

    particle_t pt = particles[tid];

    pt.vx += pt.ax * dt;
    pt.vy += pt.ay * dt;
    pt.x  += pt.vx * dt;
    pt.y  += pt.vy * dt;

    while( pt.x < 0 || pt.x > size )
    {
        pt.x  = pt.x < 0 ? -(pt.x) : 2*size-pt.x;
        pt.vx = -(pt.vx);
    }
    while( pt.y < 0 || pt.y > size )
    {
        pt.y  = pt.y < 0 ? -(pt.y) : 2*size-pt.y;
        pt.vy = -(pt.vy);
    }
    particles[tid] = pt;

}



int main( int argc, char **argv )
{    
    // This takes a few seconds to initialize the runtime
    cudaThreadSynchronize(); 

    if( find_option( argc, argv, "-h" ) >= 0 )
    {
        printf( "Options:\n" );
        printf( "-h to see this help\n" );
        printf( "-n <int> to set the number of particles\n" );
        printf( "-t <int> to set the number of threads\n");
        printf( "-o <filename> to specify the output file name\n" );
        return 0;
    }

    int n = read_int( argc, argv, "-n", 1000 );
    int num_threads = read_int(argc, argv, "-t", 256);

    char *savename = read_string( argc, argv, "-o", NULL );

    FILE *fsave = savename ? fopen( savename, "w" ) : NULL;
    particle_t *particles = (particle_t*) malloc( n * sizeof(particle_t) );

    // GPU particle data structure
    particle_t * d_particles;

    cudaError_t reterr;
    if (reterr = cudaMalloc((void **) &d_particles, n * sizeof(particle_t)))
    {
        printf("particles: Error code: %d\n", reterr);
        return -1;
    }

    set_size( n );

    // create bins: we will use this to store the index of objects
    // It is highly inefficient: we may well end up with a large vector
    int NB = (int)floor( size / cutoff) + 1; // NB is the number of bins per edge
    printf("Number of bins: %d * %d\n", NB, NB);
    unsigned int * d_bins = NULL;
    unsigned int * d_binMax = NULL;

    if (reterr = cudaMalloc((void**) & d_binMax, NB*NB*sizeof(unsigned int)))
    {
        printf("bins: Error code: %d\n", reterr);
        return -1;
    }
    if (reterr = cudaMalloc((void**) & d_bins, NB*NB*P_PER_B * sizeof(unsigned int)))
    {
        printf("binMax: Error code: %d\n", reterr);
        return -1;
    }


    init_particles( n, particles );

    cudaThreadSynchronize();
    double copy_time = read_timer( );

    // Copy the particles to the GPU
    if (reterr = cudaMemcpy(d_particles, particles, n * sizeof(particle_t), cudaMemcpyHostToDevice))
    {
        printf("copy memory error. Error code: %d\n", reterr);
        return -1;
    }

    // TODO: assign particles to bins;
    int blks = (n + num_threads - 1) / num_threads;
    printf("number of blocks: %d, number of threads: %d\n", blks, num_threads);
    cudaMemset(d_binMax, 0, NB*NB*sizeof(unsigned int));
    assign_bins_gpu <<<blks, num_threads >>> (d_particles, n, d_bins, d_binMax, NB);

    cudaThreadSynchronize();
    copy_time = read_timer( ) - copy_time;
    //
    //  simulate a number of time steps
    //
    cudaThreadSynchronize();
    double simulation_time = read_timer( );

    for( int step = 0; step < NSTEPS; step++ )
    {
        //
        //  compute forces
        //
        //printf("Iter %d\n", step);
        //printf("Compute forces\n");
        compute_forces_gpu <<< blks, num_threads >>> (d_particles, n, d_bins, d_binMax, NB);
        cudaThreadSynchronize();
        //
        //  move particles
        //
        //printf("Cuda last error: %d\n", cudaGetLastError());
        //printf("Moving particles\n");
        move_gpu <<< blks, num_threads >>> (d_particles, n, size);
        cudaThreadSynchronize();

        //
        // assign bins
        //
        //printf("Assigning bins\n");
        cudaMemset(d_binMax, 0, NB*NB*sizeof(unsigned int));
        assign_bins_gpu <<<blks, num_threads >>> (d_particles, n, d_bins, d_binMax, NB);
        cudaThreadSynchronize();

        //
        //  save if necessary
        //
        if( fsave && (step%SAVEFREQ) == 0 ) {
            // Copy the particles back to the CPU
            cudaMemcpy(particles, d_particles, n * sizeof(particle_t), cudaMemcpyDeviceToHost);
            printf("Saving!\n");
            save( fsave, n, particles);
        }
    }
    cudaThreadSynchronize();
    simulation_time = read_timer( ) - simulation_time;

    printf( "CPU-GPU copy time = %g seconds\n", copy_time);
    printf( "n = %d, simulation time = %g seconds\n", n, simulation_time );

    free( particles );
    cudaFree(d_particles);
    cudaFree(d_bins);
    cudaFree(d_binMax);

    if( fsave )
        fclose( fsave );

    return 0;
}
